package com.bsj.media.media.packageFlv;


import com.bsj.power.common.def.protocol.rtp.RtpParameterSet;
import com.bsj.power.common.util.ExceptionUtil;
import com.bsj.tools.Conv;
import org.apache.commons.lang3.tuple.Pair;

import java.nio.ByteBuffer;
import java.util.List;

/**
 * @Description flv打包工具类
 * @Date 2022/10/24 19:46
 **/
public class FlvPacketUtil {

    /**
     * 组flv的头
     *
     * @param dst:
     * @param flag: 内容标识： 0x01:视频，0x04:音频，0x05:音视频
     * @return 组装好后的偏移量
     **/
    public static void getFlvHeader(ByteBuffer dst, byte flag) {
        // 前三个byte是文件类型，总是"FLV"
        dst.put((byte) 0x46);
        dst.put((byte) 0x4C);
        dst.put((byte) 0x56);

        // 第4btye是版本号，目前一般是0x01。
        dst.put((byte) 0x01);

        dst.put(flag);

        // flv头的长度
        dst.putInt(9);
    }

    /**
     * 寻找关键帧的偏移量（00 00 01 或者00 00 00 01开始的偏移量）
     *
     * @param body:
     * @param offset:
     * @param encodingFormat: 编码格式 1-h264, 3-h265
     * @return
     **/
    public static int findIDROffset2(byte[] body, int offset, int encodingFormat) {
        int bodyLen = body.length;
        while (offset + 5 < bodyLen) {
            if (body[offset] != 0x00 || body[offset + 1] != 0x00) {
                offset++;
                continue;
            }
            if (body[offset + 2] == 0x01 || (body[offset + 2] == 0x00 && body[offset + 3] == 0x01)) {
                int idrOffset = body[offset + 2] == 0x01 ? offset + 3 : offset + 4;
                // 判断关键帧
                if (encodingFormat == RtpParameterSet.fdVideoCodeTypeH264) {
                    if ((body[idrOffset] & 0x1f) == RtpParameterSet.naluTypeIDRH4) {
                        return offset;
                    }
                } else if (encodingFormat == RtpParameterSet.fdVideoCodeTypeH265) {
                    int shortNetOrder = Conv.getShortNetOrder(body, idrOffset);
                    if (((shortNetOrder & 0x7E00) >> 9) == RtpParameterSet.naluTypeIDRH5) {
                        return offset;
                    }
                }
            }
            offset++;
        }
        return 0;
    }

    /**
     * 打包scriptTag + TagSize
     *
     * @param dst:             存放flv数据
     * @param videoFrameRate:  视频帧率
     * @param width:           视频宽
     * @param height:          视频高
     * @param videocodecid:    视频编码类型  7：H264  12:H265
     * @param audioSampleRate: 音频采样率
     * @param sampleSize:      音频采样大小
     * @param audioCodecid:    音频编码类型 10 ACC
     * @return 放回当前的tagsize
     **/
    public static int setScriptTag(ByteBuffer dst, double videoFrameRate, double width, double height, int videocodecid,
                                   int audioSampleRate, int sampleSize, int audioCodecid, int previousTagSize) {
        // put前一个tag的size
        dst.putInt(previousTagSize);
        int tagSizeStartOffset = dst.position(); // 用于计算TagSize的大小
        // position跳过11字节，先组TagData
        dst.position(tagSizeStartOffset + 11);
        // 组tagData
        int tagDataStartOffset = dst.position(); // tagData的起始偏移量

        // ===========================
        setScriptTagBody(dst, videoFrameRate, width, height, videocodecid, audioSampleRate, sampleSize, audioCodecid);
        // ===========================

        int tagSizeEndOffset = dst.position();
        // 组TagHeader
        setTagHeader((byte) 18, tagSizeEndOffset - tagDataStartOffset, 0, tagSizeStartOffset, dst);
        return tagSizeEndOffset - tagSizeStartOffset;
    }

    /**
     * 打包H264的flv第一视频帧
     *
     * @param dst:
     * @param sps:
     * @param pps:
     * @return 返回当前的TagSize
     **/
    public static int setH264FirstVideoTag(ByteBuffer dst, byte[] sps, byte[] pps, int previousTagSize) {
        dst.putInt(previousTagSize); // put前一个tag的size
        int tagSizeStartOffset = dst.position(); // 用于计算TagSize的大小
        // position跳过11字节，先组TagData
        dst.position(tagSizeStartOffset + 11);
        // 组tagData
        int tagDataStartOffset = dst.position(); // tagData的起始偏移量

        //================================
        dst.put((byte) (0x10 | RtpParameterSet.flvCodecIdH264)); // Frametype and CodecID
        dst.put((byte) 0); // AVC sequence header
        intPut3Bytes(dst, 0); // composition time
        dst.put((byte) 1); // version

        // sps数据可能是 00 00 01 0x67开头，所以需要判断帧头长度
        int spsHeaderLength = sps[2] == 1 ? 3 : 4;
        dst.put(sps[1 + spsHeaderLength]);
        dst.put(sps[2 + spsHeaderLength]);
        dst.put(sps[3 + spsHeaderLength]);

        dst.put((byte) 0xff);
        dst.put((byte) 0xe1);

        // sps长度 + 数据
        dst.putShort((short) (sps.length - spsHeaderLength));
        dst.put(sps, spsHeaderLength, sps.length - spsHeaderLength);

        dst.put((byte) 0x01); // pps数量

        // pps长度+数据
        // pps数据可能是 00 00 01 0x68开头，所以需要判断帧头长度
        int ppsHeaderLength = pps[2] == 1 ? 3 : 4;
        dst.putShort((short) (pps.length - ppsHeaderLength));
        dst.put(pps, ppsHeaderLength, pps.length - ppsHeaderLength);
        //================================

        int tagSizeEndOffset = dst.position();
        // 组TagHeader
        setTagHeader((byte) 9, tagSizeEndOffset - tagDataStartOffset, 0, tagSizeStartOffset, dst);
        return tagSizeEndOffset - tagSizeStartOffset;
    }

    /**
     * 打包H265的flv第一视频帧
     *
     * @param dst:
     * @param vps:
     * @param sps:
     * @param pps:
     * @return
     **/
    public static int setH265FirstVideoTag(ByteBuffer dst, byte[] vps, byte[] sps, byte[] pps, int previousTagSize) {
        dst.putInt(previousTagSize); // put前一个tag的size
        int tagSizeStartOffset = dst.position(); // 用于计算TagSize的大小
        // position跳过11字节，先组TagData
        dst.position(tagSizeStartOffset + 11);
        // 组tagData
        int tagDataStartOffset = dst.position(); // tagData的起始偏移量

        //======================================
        // Frametype and CodecID
        dst.put((byte) 0x1C);
        dst.putInt(0);

        // Hevc
        dst.put((byte) 0x01); // configurationVersion = 01
        dst.put((byte) 0x01); // general_profile_space(2) + general_tier_flag(1) + general_profile_idc(5)
        byteBufferSetNbitToZero(dst, 19);// 跳过19个字节 相当于 dst.put(new byte[20]);
        //dst.position(dst.position() + 20); 这里直接跳过，原本的位置会残留之前组装的数据

        // 对vps/sps/pps封装
        dst.put((byte) 0x03); // lengthSizeMinusOne
        dst.put((byte) 0x03); // numOfArrays
        // vps
        dst.put((byte) 0x20); // vps 32
        dst.put((byte) ((1 >> 8) & 0xff));
        dst.put((byte) 1);
        int vpsHeadLen = (vps[2] == 0x01) ? 3 : 4;
        dst.putShort((short) (vps.length - vpsHeadLen));
        dst.put(vps, vpsHeadLen, vps.length - vpsHeadLen);

        // sps
        int spsHeadLen = (sps[2] == 0x01) ? 3 : 4; // sps帧头大小
        dst.put((byte) 0x21); // sps 33
        dst.put((byte) ((1 >> 8) & 0xff));
        dst.put((byte) 1);
        dst.putShort((short) (sps.length - spsHeadLen));
        dst.put(sps, spsHeadLen, sps.length - spsHeadLen);

        // pps
        dst.put((byte) 0x22); // pps 34
        dst.put((byte) ((1 >> 8) & 0xff));
        dst.put((byte) 1);
        int ppsHeadLen = (pps[2] == 0x01) ? 3 : 4;
        dst.putShort((short) (pps.length - ppsHeadLen));
        dst.put(pps, ppsHeadLen, pps.length - ppsHeadLen);
        //======================================

        int tagSizeEndOffset = dst.position();
        // 组TagHeader
        setTagHeader((byte) 9, tagSizeEndOffset - tagDataStartOffset, 0, tagSizeStartOffset, dst);
        return tagSizeEndOffset - tagSizeStartOffset;
    }

    /**
     * 打包第一flv音频包
     *
     * @param dst:
     * @param aacSampleRate:
     * @param aacChannel:
     * @return
     **/
    public static int setAACFirstTag(ByteBuffer dst, int aacSampleRate, int aacChannel, int audiocodecid, int previousTagSize) {
        dst.putInt(previousTagSize); // put前一个tag的size
        int tagSizeStartOffset = dst.position(); // 用于计算TagSize的大小
        // position跳过11字节，先组TagData
        dst.position(tagSizeStartOffset + 11);
        // 组tagData
        int tagDataStartOffset = dst.position(); // tagData的起始偏移量

        // ====================================
        dst.put((byte) ((audiocodecid << 4) | (aacSampleRate << 2) | (1 << 1) | aacChannel));  //标识SoundFomat SoundSize SoundRate SoundFormat
        dst.put((byte) 0x00);

        byte temp1 = 0x2 << 3; // Object Type(5bits) 2 -AAC-LC 5 -SBR 29 -PS
        // SampleRate Index(4bits):0 -96000, 1-88200, 3-64000, 4-44100 5-32000 6- 24000, 7 -22050 8 -16000, 11-8000
        if (aacSampleRate == 3) {// 44k
            temp1 |= (0x4 >> 1);
        } else if (aacSampleRate == 1) {// 11K--> 16000
            temp1 |= (0x08 >> 1);
        }
        dst.put(temp1);
        dst.put((byte) (0x1 << 3)); // 单声道
        dst.put(new byte[3]); //加三个字节  禁止这样赋值：dst.position(dst.position() + 3); 因为会残留之前数据
        // ====================================

        int tagSizeEndOffset = dst.position();
        //dst.putInt(tagSizeEndOffset - tagSizeStartOffset); // TagSize
        // 组TagHeader
        setTagHeader((byte) 8, tagSizeEndOffset - tagDataStartOffset, 0, tagSizeStartOffset, dst);
        return tagSizeEndOffset - tagSizeStartOffset;
    }

    /**
     * 打包flv的视频Tag
     *
     * @param body:
     * @param startOffset:
     * @param len:
     * @param dst:
     * @param codecid:
     * @param keyFrame:
     * @param timeStamp:
     * @param frameRate:
     * @return
     **/
    public static int setVideoTag(byte[] body, int startOffset, int len, ByteBuffer dst, byte codecid, int keyFrame, int timeStamp, int frameRate, int previousTagSize) {
        dst.putInt(previousTagSize); // put前一个tag的size
        int tagSizeStartOffset = dst.position(); // 用于计算TagSize的大小
        // position跳过11字节，先组TagData
        dst.position(tagSizeStartOffset + 11);
        // 组tagData
        int tagDataStartOffset = dst.position(); // tagData的起始偏移量

        //============================================
        // 高4比特： 1、关键帧； 2、内部帧（非IDR帧），  低4比特： 7：H264 12:H265
        dst.put((byte) (keyFrame << 4 | codecid));
        dst.put((byte) 0x01); // AVCPacketType
        intPut3Bytes(dst, frameRate); // 时间的偏移量
        dst.putInt(len); // 帧大小
        dst.put(body, startOffset, len); // 数据区
        //============================================
        int tagSizeEndOffset = dst.position();

        // 组TagHeader
        setTagHeader((byte) 9, tagSizeEndOffset - tagDataStartOffset, timeStamp, tagSizeStartOffset, dst);
        return tagSizeEndOffset - tagSizeStartOffset;
    }

    /**
     * 打包flv的视频Tag
     *
     * @param body:
     * @param dst:
     * @param codecid:
     * @param keyFrame:
     * @param timeStamp:
     * @param frameRate:
     * @param copyInfo:  需要打包视频帧，body的偏移量信息， 里面元素Pair<Integer, Integer>表示一个帧信息，left:起始偏移量，right:需要赋值长度
     * @return
     **/
    public static int setVideoTag(byte[] body, ByteBuffer dst, byte codecid, int keyFrame, int timeStamp, int frameRate, int previousTagSize, List<Pair<Integer, Integer>> copyInfo) {
        dst.putInt(previousTagSize); // put前一个tag的size
        int tagSizeStartOffset = dst.position(); // 用于计算TagSize的大小
        // position跳过11字节，先组TagData
        dst.position(tagSizeStartOffset + 11);
        // 组tagData
        int tagDataStartOffset = dst.position(); // tagData的起始偏移量

        //============================================
        // 高4比特： 1、关键帧； 2、内部帧（非IDR帧），  低4比特： 7：H264 12:H265
        dst.put((byte) (keyFrame << 4 | codecid));
        dst.put((byte) 0x01); // AVCPacketType
        intPut3Bytes(dst, frameRate); // 时间的偏移量
        copyInfo.forEach(E -> {
            dst.putInt(E.getRight()); // 帧大小
            dst.put(body, E.getLeft(), E.getRight()); // 数据区
        });
        //============================================
        int tagSizeEndOffset = dst.position();

        // 组TagHeader
        setTagHeader((byte) 9, tagSizeEndOffset - tagDataStartOffset, timeStamp, tagSizeStartOffset, dst);
        return tagSizeEndOffset - tagSizeStartOffset;
    }

    /**
     * @param aac:           aac原始数据
     * @param dst:
     * @param codecid:       编码id 10:ACC
     * @param aacSampleRate: 采样率
     * @param aacChannel:    音频声道数
     * @param timeStamp:
     * @return
     **/
    public static int setAudioTag(byte[] aac, ByteBuffer dst, int codecid, int aacSampleRate, int aacChannel, int timeStamp, int previousTagSize) {
        dst.putInt(previousTagSize); // put前一个tag的大小
        int tagSizeStartOffset = dst.position(); // 用于计算TagSize的大小
        // position跳过11字节，先组TagData
        dst.position(tagSizeStartOffset + 11);
        // 组tagData
        int tagDataStartOffset = dst.position(); // tagData的起始偏移量

        //===========================================
        dst.put((byte) ((codecid << 4) | (aacSampleRate << 2) | (1 << 1) | aacChannel));  //标识SoundFomat SoundSize SoundRate SoundFormat
        dst.put((byte) 0x01);
        dst.put(aac, 0, aac.length); // 音频数据
        //===========================================

        int tagSizeEndOffset = dst.position();
        //dst.putInt(tagSizeEndOffset - tagSizeStartOffset); // TagSize
        // 组TagHeader
        setTagHeader((byte) 8, tagSizeEndOffset - tagDataStartOffset, timeStamp, tagSizeStartOffset, dst);
        return tagSizeEndOffset - tagSizeStartOffset;
    }

    /**
     * 打包ScriptTagBody
     **/
    private static void setScriptTagBody(ByteBuffer dst, double videoFrameRate, double width, double height, int videocodecid, int audioSampleRate, int sampleSize, int audioCodecid) {
        // 第一个AMF包
        dst.put((byte) 0x02); // 字符串类型
        setAmfString(dst, "onMetaData");
        // 第二个AMF包
        dst.put((byte) 0x08); // ECMA array type类型
        int paramSizeOffset = dst.position();
        dst.position(paramSizeOffset + 4); // 跳过参数个数
        int paramSize = 0;
        // 视频
        if (width != 0 && height != 0 && videoFrameRate != 0) {
            setAmfString(dst, "width");
            setAmfDouble(dst, width);

            setAmfString(dst, "height");
            setAmfDouble(dst, height);

            setAmfString(dst, "framerate");
            setAmfDouble(dst, videoFrameRate);
            // 视频编码方式
            setAmfString(dst, "videocodecid");
            setAmfDouble(dst, videocodecid);
            paramSize += 4;
        }

        // 音频
        if (audioSampleRate != 0 && audioCodecid != 0 && sampleSize != 0) {
            // 音频采样率
            setAmfString(dst, "audiosamplerate");
            setAmfDouble(dst, audioSampleRate);

            // 音频精度
            setAmfString(dst, "audiosamplesize");
            setAmfDouble(dst, sampleSize);

            // 音频是否为立体声:1立体声 0:不是立体声
            setAmfString(dst, "stereo");
            setAmfBool(dst, (byte) 0);

            // 音频编码方式
            setAmfString(dst, "audiocodecid");
            setAmfDouble(dst, audioCodecid);
            paramSize += 4;
        }

        // 设备参数个数
        dst.putInt(paramSizeOffset, paramSize);

        // 结束符
        setAmfString(dst, "");
        dst.put((byte) 0x09);
    }

    /**
     * 对TagHeader打包， tagData已经在dst中，只是对tagData前的11个字节进行填充
     *
     * @param tagType:     8为Audio，9为Video，18为Scripts
     * @param dataSize:    tagData大小
     * @param timeStamp:   时间戳
     * @param startOffset: dst的操作偏移量
     * @param dst:
     * @return
     **/
    public static void setTagHeader(byte tagType, int dataSize, int timeStamp, int startOffset, ByteBuffer dst) {
        // tagType
        dst.put(startOffset++, tagType);

        // tagData大小
        dst.put(startOffset++, (byte) ((dataSize & 0x00ff0000) >> 16));
        dst.put(startOffset++, (byte) ((dataSize & 0x0000ff00) >> 8));
        dst.put(startOffset++, (byte) ((dataSize & 0x000000ff) >> 0));

        // 时间戳
        dst.put(startOffset++, (byte) ((timeStamp & 0x00ff0000) >> 16));
        dst.put(startOffset++, (byte) ((timeStamp & 0x0000ff00) >> 8));
        dst.put(startOffset++, (byte) ((timeStamp & 0x000000ff) >> 0));

        // 时间戳后面一个byte是扩展时间戳，时间戳不够长的时候用
        dst.put(startOffset++, (byte) (timeStamp >> 24 & 0xFF));

        // 最后三个是streamID，但是总为0。所以不进行操作
        dst.put(startOffset++, (byte) 0);
        dst.put(startOffset++, (byte) 0);
        dst.put(startOffset++, (byte) 0);
    }

    /**
     * 把值put到ByteBuffer
     *
     * @param dst:
     * @param value:
     * @return
     **/
    public static void intPut3Bytes(ByteBuffer dst, int value) {
        dst.put((byte) ((value & 0x00ff0000) >> 16));
        dst.put((byte) ((value & 0x0000ff00) >> 8));
        dst.put((byte) ((value & 0x000000ff) >> 0));
    }

    /**
     * amf的String打包方式
     *
     * @param dst:
     * @param str:
     * @return
     **/
    static void setAmfString(ByteBuffer dst, String str) {
        // 2字节的长度
        dst.putShort((short) str.length());
        // 后面是字符串内容
        dst.put(str.getBytes());
    }

    static void setAmfDouble(ByteBuffer dst, double value) {
        dst.put((byte) 0);
        dst.putLong(Double.doubleToRawLongBits(value));
    }

    static void setAmfBool(ByteBuffer dst, byte b) {
        dst.put((byte) 1);
        dst.put(b);
    }

    /**
     * 寻找关键帧的偏移量（00 00 01 或者00 00 00 01后的偏移量）
     *
     * @param body:
     * @param offset:
     * @param encodingFormat: 编码格式 1-h264, 3-h265
     * @return
     **/
    public static int findIDROffset(byte[] body, int offset, int encodingFormat) {
        int bodyLen = body.length;
        while (offset + 5 < bodyLen) {
            if (body[offset] != 0x00 || body[offset + 1] != 0x00) {
                offset++;
                continue;
            }
            if (body[offset + 2] == 0x01 || (body[offset + 2] == 0x00 && body[offset + 3] == 0x01)) {
                int idrOffset = body[offset + 2] == 0x01 ? offset + 3 : offset + 4;
                // 判断关键帧
                if (encodingFormat == RtpParameterSet.fdVideoCodeTypeH264) {
                    if ((body[idrOffset] & 0x1f) == RtpParameterSet.naluTypeIDRH4) {
                        return idrOffset;
                    }
                } else if (encodingFormat == RtpParameterSet.fdVideoCodeTypeH265) {
                    int shortNetOrder = Conv.getShortNetOrder(body, idrOffset);
                    if (((shortNetOrder & 0x7E00) >> 9) == RtpParameterSet.naluTypeIDRH5) {
                        return idrOffset;
                    }
                }
            }
            offset++;
        }
        return 0;
    }

    /**
     * 设置dst当前偏移量之后的len位，都为0
     *
     * @param dst:
     * @param len:
     * @return
     **/
    public static void byteBufferSetNbitToZero(ByteBuffer dst, int len) {
        if (dst == null) {
            return;
        }
        try {
            for (int i = 0; i < len; i++) {
                dst.put((byte) 0);
            }
        } catch (Exception e) {
            System.out.println("ByteBuffer处理失败：" + ExceptionUtil.getStackStr(e));
        }
    }
}
